import os
from PIL import Image
import random


class data_split():

    def __init__(self, PATH):
        self.path = PATH

    # 判断是否有文件
    def isfile(self):
        for folder in os.listdir(self.path):
            if not os.path.isdir(self.path+folder):
                os.remove(self.path+folder)

    # 建立文件夹
    def mkdir(self, file_path):
        isexists = os.path.exists(self.path)
        if not isexists:
            os.makedirs(self.path)
            print(self.path)
            print('success')
            return True
        else:
            print(self.path)
            print('folder is exist')
            return False

    # 遍历每个文件夹
    def eachFile(self, file):
        pathDir = os.listdir(self.path)
        child_file_name = []
        full_child_file_list = []
        for allDir in pathDir:
            if not allDir == '.DS_Store':
                child = os.path.join(self.path, allDir)

                full_child_file_list.append(child)
                child_file_name.append(allDir)
        return full_child_file_list, child_file_name

    # 验证集与训练集的分配率
    def move_ratio(self, data_list, original_str, replace_str):
        for x in data_list:
            fromImage = Image.open(x)
            x = x.replace(original_str, replace_str)
            fromImage.save(x)


if __name__ == '__main__':
    data_path = './data/train/'
    data_tra_path = './data/train_data/'
    data_val_path = './data/val_data/'

    train_val = data_split(PATH=data_path)
    full_child_file, child_file = train_val.eachFile(data_path)

    for i in child_file:
        tra_path = data_tra_path + '/' + str(i)
        train_val.mkdir(tra_path)
        val_path = data_val_path + '/' + str(i)
        train_val.mkdir(val_path)

    test_train_split_ratio = 0.85

    for i in full_child_file:
        pic_dir, pic_name = train_val.eachFile(i)
        random.shuffle(pic_dir)
        train_list = pic_dir[0:int(test_train_split_ratio * len(pic_dir))]
        val_list = pic_dir[int(test_train_split_ratio * len(pic_dir)):]
        # train_move, val_move
        print('proprecessing %s' % i)

        train_val.move_ratio(train_list, data_path, data_tra_path)
        train_val.move_ratio(val_list, data_path, data_val_path)